import os
import shutil
from ftplib import FTP

from config import Config


class FTPManager:
    def __init__(self):
        self.config = Config()
        self.ftp = FTP()

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.disconnect()

    def connect(self):
        self.ftp.connect(self.config.FTP_HOST, self.config.FTP_PORT)
        self.ftp.login(self.config.FTP_USER, self.config.FTP_PASSWORD)

    def disconnect(self):
        self.ftp.quit()

    def upload_userdata(self, user_name: str, local_dir: str):
        local_dir = os.path.abspath(local_dir)

        def upload_dir(root_dir):
            remote_dir = os.path.abspath(root_dir).replace(local_dir, f'~/{user_name}/', 1).rstrip("/")
            self.ftp.cwd("/".join(remote_dir.split("/")[:-1]))
            try:
                self.ftp.mkd(remote_dir)
            except Exception:
                pass
            self.ftp.cwd(os.path.basename(remote_dir))
            dir_list = []
            for file in os.listdir(root_dir):
                local_file = os.path.join(root_dir, file)
                if os.path.isdir(local_file):
                    dir_list.append(local_file)
                    continue
                if os.path.isfile(local_file):
                    with open(local_file, 'rb') as f:
                        self.ftp.storbinary(f'STOR {file}', f)

            for sub_dir in dir_list:
                upload_dir(sub_dir)

        upload_dir(local_dir)

    def download_userdata(self, user_name: str, local_dir: str):
        local_dir = os.path.abspath(local_dir)
        remote_root_dir = f"{user_name}/crypto"

        def download_dir(remote_dir):
            current_local_dir = os.path.join(local_dir, remote_dir[len(remote_root_dir):].lstrip("/").rstrip("/"))
            if os.path.exists(current_local_dir):
                if os.path.isdir(current_local_dir):
                    shutil.rmtree(current_local_dir)
                else:
                    os.remove(current_local_dir)
            if not os.path.exists(current_local_dir):
                os.makedirs(current_local_dir, exist_ok=True)

            self.ftp.cwd("~/" + remote_dir)

            dir_list = []
            for file in self.ftp.nlst():
                try:
                    local_file = os.path.join(current_local_dir, file)
                    with open(local_file, 'wb') as f:
                        self.ftp.retrbinary(f'RETR {file}', f.write)
                except Exception as e:
                    print(e)
                    dir_list.append(os.path.join(remote_dir, file))

            for sub_dir in dir_list:
                try:
                    download_dir(sub_dir)
                except Exception as e:
                    print(e)

        download_dir(remote_root_dir)

    def remove_userdata(self, user_name):
        try:
            self.ftp.cwd("~/")
            self.ftp.rmd(user_name)
        except Exception:
            pass

    def list_users(self, tenant_id):
        self.connect()
        tenant_dir = self.ensure_tenant_dir(tenant_id)

        try:
            files = []
            self.ftp.dir(tenant_dir, files.append)
            user_files = [f.split()[-1] for f in files if f.startswith('-') and f.endswith('.json')]
            self.disconnect()
            return user_files
        except:
            self.disconnect()
            return []
